from dataset import Data, Data3
from sklearn import tree, linear_model, svm, neighbors, ensemble
from sklearn.metrics import mean_absolute_percentage_error, mean_absolute_error, mean_squared_error
from visual import plot, dfplot, plot2
import pandas as pd
import numpy as np
from sklearn.metrics import r2_score
from sklearn.model_selection import train_test_split

test_size = 0.7

class SVMModel:
    def __init__(self, args):
        self.args = args

        self._init_data()
        
    def _init_data(self):
        data = Data3(self.args.data_path)
        self.data1_date, self.data1_X, self.data1_y = data.data1
        self.data2_date, self.data2_X, self.data2_y = data.data2
        self.train_X1, _, self.train_y1, _ = train_test_split(self.data1_X, self.data1_y, test_size=test_size)
        self.train_X2, _, self.train_y2, _ = train_test_split(self.data2_X, self.data2_y, test_size=test_size)


    def train_dt(self):
        self.model1 = svm.SVR()
        self.model2 = svm.SVR()
        self.model1.fit(self.train_X1, self.train_y1)
        self.model2.fit(self.train_X2, self.train_y2)

    def test(self):
        print(self.test_X.shape)
        pred = self.model.predict(self.test_X)
        mae = mean_absolute_percentage_error(self.test_y, pred)
        print("mae_loss: %.4f" % (mae))

        df = pd.DataFrame()
        df['date'] = self.test_date
        df['y'] = self.test_y
        df['pred'] = pred
        new_df = df.groupby('date').mean()
        # print(new_df)
        plot(new_df.index.tolist(), [new_df.y, new_df.pred], ['真实值', '预测值'], 'date', 'value', './results/')
        
    def test2(self):
        """
        散点图
        """
        
        pred1 = self.model1.predict(self.data1_X)
        pred2 = self.model2.predict(self.data2_X)

        mse1 = np.sqrt(np.mean((self.data1_y - pred1)**2))      # 均方根误差
        mae1 = np.mean(np.abs(self.data1_y - pred1))            # 平均绝对值误差
        r2_1 = r2_score(self.data1_y, pred1)
        print("均方根误差csv: %.4f" % (mse1))  
        print("平均绝对误差csv: %.4f" % mae1)
        print("r2_csv", r2_1)
        
        mse2 = np.sqrt(np.mean((self.data2_y - pred2)**2))      # 均方根误差
        mae2 = np.mean(np.abs(self.data2_y - pred2))            # 平均绝对值误差
        r2_2 = r2_score(self.data2_y, pred2)
        print("均方根误差excel: %.4f" % (mse2))  
        print("平均绝对误差excel: %.4f" % mae2)
        print("r2_csv", r2_2)
        print()
        # print(pred.shape, self.test_y.shape)
        # df = pd.DataFrame()
        # df['date'] = self.test_date
        # df['y'] = self.test_y
        # df['pred'] = pred
        n_csv = len(set(self.data1_y))
        date = pd.concat([self.data1_date, self.data2_date], axis=0, ignore_index=True)
        pred = np.concatenate([pred1, pred2])
        real = np.concatenate([self.data1_y, self.data2_y])

        df = pd.DataFrame({'date': date, 'pred':pred, 'real':real})
        # df = df.drop_duplicates(subset=['date'], keep='first')
        # print((df.pred-pred2).mean())
        plot2(df.real, df.pred, n=2880, xlabel="Actual Value", ylabel="Estimate", save_path="./results/SVM", filename='总的')
        plot2(self.data1_y, pred1, n=n_csv, xlabel="Actual Value", ylabel="Estimate", save_path="./results/SVM", filename='csv')
        plot2(self.data2_y, pred2, n=2880-n_csv, xlabel="Actual Value", ylabel="Estimate", save_path="./results/SVM", filename='excel')
        df = df.groupby('date').mean()
        plot(df.index.tolist(), [df.real, df.pred], ['Actual Value', 'Estimate'], 'date', 'Visibility (m)', './results/SVM')

